/*    Copyright 2010 Tobias Marschall
 *
 *    This file is part of MoSDi.
 *
 *    MoSDi is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    MoSDi is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with MoSDi.  If not, see <http://www.gnu.org/licenses/>.
 */

package mosdi.fa;

import java.util.Arrays;

import mosdi.discovery.TextModelIupacBounds;
import mosdi.util.BitArray;
import mosdi.util.Iupac;
import mosdi.util.IupacStringConstraints;

/** Class able to compute a pattern's expectation given a text model. This class is 
 *  designed to recycle intermediate results when the pattern is changed. */
public class ForwardPatternExpectationCalculator implements PatternExpectationCalculator {
	private static final BitArray[] generalizedAlphabet = Iupac.asGeneralizedAlphabet();
	private FiniteMemoryTextModel textModel;
	/**
	 * position within distribution (with respect to the pattern-length) where the distribution is still valid
	 * validUpTo = 0: position pattern[0] is calculated within distributions
	 * distributions[validUpTo][x] has to be valid
	 */
	private int validUpTo;
	private int[] pattern;
	private double[][] distributions;
	/**
	 * stores the sum of the distribution. 
	 * sum_i=0^t { sum_j=0^stateCount { distribution[i][j] } } = patternExpectation[t]
	 */
	private double[] patternExpectation;
	private TextModelIupacBounds textModelBounds;

	/** Constructor.
	 * 
	 * @param pattern Pattern with respect to the IUPAC alphabet.
	 */
	public ForwardPatternExpectationCalculator(FiniteMemoryTextModel textModel, int[] pattern) {
		this(textModel, pattern, new TextModelIupacBounds(textModel,false));
	}
	
	
	/** Constructor.
	 * 
	 * @param pattern Pattern with respect to generalized alphabet.
	 */
	public ForwardPatternExpectationCalculator(FiniteMemoryTextModel textModel, int[] pattern, TextModelIupacBounds textModelBounds) {
		this.textModelBounds = textModelBounds;
		// check whether each 'character' of the generalizedAlphabet has the length 
		// of the alphabet
		for(int i=0; i<generalizedAlphabet.length; ++i) {
			if(generalizedAlphabet[i].size() != textModel.getAlphabetSize()) {
				throw new IllegalArgumentException("Alphabet size mismatch");
			}
		}
		this.textModel = textModel;
		this.pattern = pattern;
		this.distributions = new double[pattern.length][textModel.getStateCount()];
		this.patternExpectation = new double[pattern.length]; 
		validUpTo = -1;
	}

	/** Returns the expectation of the patterns prefix. */
	public double getPrefixExpectation(int prefixLength) {
		updatePrefixDistribution(prefixLength-1);
		return (prefixLength>0)?patternExpectation[prefixLength-1]:1.0;
	}
	
	/**
	 * 
	 * @param position Index in pattern[0,...,n-1] where charakter pattern[position] shall be changed
	 * @param newChar the new character which shall replace pattern[position]
	 */
	public void setPatternPosition(int position, int newChar) {
		if(position < 0 || position >= pattern.length) {
			throw new IllegalArgumentException("position ("+position+") is not valid. Choose within the range [0,"+(pattern.length-1)+"]");
		}
		if (pattern[position]==newChar) return;
		validUpTo = Math.min(validUpTo,position-1); 
		pattern[position] = newChar;
	}

	/**
	 * @param patternPosition the prefix pattern[0,..., patternPosition] shall be updated within distributions
	 */
	private void updatePrefixDistribution(int patternPosition) {
		double[] lastDistribution = null;
		while (validUpTo < patternPosition) { // be lazy
			if (validUpTo < 0) {
				lastDistribution = textModel.getEquilibriumDistribution();
			} else {
				lastDistribution = distributions[validUpTo];
			}
			Arrays.fill(distributions[validUpTo + 1], 0.0d);

			for (int state = 0; state < textModel.getStateCount(); ++state) {
				for (int c = 0; c < textModel.getAlphabetSize(); ++c) {
					if (generalizedAlphabet[pattern[validUpTo + 1]].get(c)) {
						for (int targetState : textModel.getTransitionTargets(state,c)) {
							distributions[validUpTo + 1][targetState] += lastDistribution[state] * textModel.getProbability(state, c, targetState);
						}
					}
				}
			}
			// update the patternExpectation-Array
			patternExpectation[validUpTo+1] = 0.0d;
			for (int state = 0; state < textModel.getStateCount(); ++state) {
				patternExpectation[validUpTo+1] += distributions[validUpTo+1][state];
			} 
			++validUpTo;
		}
	}

	@Override
	public double getExpectation() {
		return getPrefixExpectation(pattern.length);
	}

	@Override
	public double getLowerBound(int prefixLength, IupacStringConstraints suffixConstraints) {
		return getExpectationLowerOrUpperBound(prefixLength, suffixConstraints, false);
	}

	@Override
	public double getUpperBound(int prefixLength, IupacStringConstraints suffixConstraints) {
		return getExpectationLowerOrUpperBound(prefixLength, suffixConstraints, true);
	}
	
	/** Common implementation of getExpectationLowerBound and getExpectationUpperBound*/
	private double getExpectationLowerOrUpperBound(int prefixLength, IupacStringConstraints suffixConstraints, boolean upper) {
		double p = getPrefixExpectation(prefixLength);
		int[] minFrequencies = suffixConstraints.getMinFrequencies();
		int[] maxFrequencies = suffixConstraints.getMaxFrequencies();
		int n = pattern.length - prefixLength;
		for (int i=0; i<4; ++i) {
			if (minFrequencies[i]==0) continue;
			p*=Math.pow(upper?textModelBounds.characterClassMaxProbabilityNoFuture(i):textModelBounds.characterClassMinProbabilityNoFuture(i),minFrequencies[i]);
			n-=minFrequencies[i];
		}
		if (upper) {
			for (int i=3; (i>=0)&&(n>0); --i) {
				if (maxFrequencies[i]==0) continue;
				int k = Math.min(maxFrequencies[i], n);
				p*=Math.pow(textModelBounds.characterClassMaxProbabilityNoFuture(i),k);
				n-=k;
			}		
		} else {
			for (int i=0; (i<4)&&(n>0); ++i) {
				if (maxFrequencies[i]==0) continue;
				int k = Math.min(maxFrequencies[i], n);
				p*=Math.pow(textModelBounds.characterClassMinProbabilityNoFuture(i),k);
				n-=k;
			}		
		}
		if (n!=0) throw new IllegalStateException();
		return p;
	}

	
}
